#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于客户提供的main.py方法的改进仿真函数
实现了opt_core3函数，支持单次仿真和批量仿真
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import logging
import os
from datetime import datetime
from opt_alg.bi_mp_pro2 import parallel_optimize
from opt_alg.param_ident2 import share_data

# 设置matplotlib中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Times New Roman']
plt.rcParams['axes.unicode_minus'] = False

def setup_logger():
    """设置日志"""
    logger = logging.getLogger("opt_core3_simulation")
    logger.setLevel(logging.INFO)
    
    # 创建控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    
    # 创建格式器
    formatter = logging.Formatter(
        '[%(asctime)s-%(name)s] (%(filename)s %(lineno)d): %(levelname)s %(message)s',
        datefmt="%Y-%m-%d %H:%M:%S"
    )
    console_handler.setFormatter(formatter)
    
    # 添加处理器到logger
    logger.addHandler(console_handler)
    
    return logger

def generate_sample_data():
    """
    基于客户提供的main.py中的数据格式生成样本数据
    返回与客户数据格式完全一致的数据结构
    """
    # 基于客户数据的时间戳
    timestamp = pd.Timestamp('2025-05-29 07:29:36')
    
    # 基于客户数据的24维参数（分析客户数据格式）
    # 客户数据: [timestamp, 469.706726, 431.0983616666666, 539.432434, 16.741943, 0.7, 2.938889, 528.03302, 60.735016, 0.0, 0.0, 442.82057027777773, 430.30466722222224, 63.234436, 91.30099166071429, 153.47908842857143, 56.0, 0.0, 18.877808, 29.437347, 0.255665, 0.0, 0.0, 3.6824282681748555, 466.02429773182513]
    
    # 解析客户数据格式（基于代码中的input_points分析）
    sample_data = [
        timestamp,                    # 0: 时间戳
        469.706726,                  # 1: 机组负荷 (MW)
        431.0983616666666,           # 2: 主汽流量 (t/h)
        539.432434,                  # 3: 主汽温度 (°C)
        16.741943,                   # 4: 主汽压力 (MPa)
        0.7,                         # 5: 再热蒸汽压力 (MPa)
        2.938889,                    # 6: 再热蒸汽温度 (°C)
        528.03302,                   # 7: 凝结水温度 (°C)
        60.735016,                   # 8: 凝结水流量 (t/h)
        0.0,                         # 9: 主给水流量 (t/h)
        0.0,                         # 10: 主给水温度 (°C)
        442.82057027777773,          # 11: 凝结水流量 (重复)
        430.30466722222224,          # 12: 凝结水泵出口温度 (°C)
        63.234436,                   # 13: 风机转速 (rpm)
        91.30099166071429,           # 14: 风机台数
        153.47908842857143,          # 15: 背压 (kPa)
        56.0,                        # 16: 风机台数 (重复)
        0.0,                         # 17: 背压 (重复)
        18.877808,                   # 18: 环境温度 (°C)
        29.437347,                   # 19: 环境温度 (重复)
        0.255665,                    # 20: 第一级抽汽压力
        0.0,                         # 21: 第二级抽汽压力
        0.0,                         # 22: 第三级抽汽压力
        3.6824282681748555,          # 23: 第四级抽汽压力
        466.02429773182513           # 24: 额外参数
    ]
    
    return sample_data

def generate_varied_data(base_data, variation_factor=0.1):
    """
    基于基础数据生成变化的数据，模拟不同运行工况
    """
    varied_data = base_data.copy()
    
    # 对关键参数添加随机变化
    key_indices = [1, 2, 3, 4, 5, 6, 7, 8, 13, 15, 18]  # 关键参数索引
    
    for idx in key_indices:
        if idx < len(varied_data) and isinstance(varied_data[idx], (int, float)) and varied_data[idx] != 0:
            # 添加±variation_factor的随机变化
            variation = np.random.uniform(-variation_factor, variation_factor)
            varied_data[idx] = varied_data[idx] * (1 + variation)
    
    return varied_data

def opt_core3(need_time, latest_one_new, latest_n_new, enough_his, logger=None):
    """
    基于客户提供的main.py方法的优化核心函数
    
    参数:
    - need_time: 需要优化的时间戳列表
    - latest_one_new: 最新一条数据
    - latest_n_new: 最新n条历史数据
    - enough_his: 历史数据是否足够
    - logger: 日志对象
    """
    
    if logger is None:
        logger = setup_logger()
    
    logger.info("开始运行opt_core3优化算法...")
    logger.info(f"优化时间: {need_time}")
    logger.info(f"历史数据足够: {enough_his}")
    
    # 检查数据有效性
    if not enough_his:
        logger.warning("历史数据不足，无法进行优化")
        return None
    
    try:
        # 获取模型参数
        ita_current, kc_current, sim_steps = share_data()
        logger.info(f"使用模型参数: 汽轮机效率={ita_current:.4f}, 凝汽器系数={kc_current:.4f}")
        
        # 处理历史数据格式
        if isinstance(latest_n_new, list) and len(latest_n_new) > 0:
            # 转换为numpy数组格式
            if isinstance(latest_n_new[0], list):
                # 已经是列表格式，直接使用
                history_data = np.array(latest_n_new)
            else:
                # 需要重新组织数据格式
                history_data = np.array([latest_n_new])
        else:
            logger.error("历史数据格式错误")
            return None
        
        logger.info(f"历史数据维度: {history_data.shape}")
        
        # 运行优化算法
        matrix_best, trace_list = parallel_optimize(
            data=history_data,
            simtime=25,
            sample_time=5,
            h=0.1,
            tolerance=0.0005,
            precision=6,
            logger=logger
        )
        
        # 解析优化结果
        if matrix_best is not None and len(matrix_best) > 0:
            optimal_fan_speed = matrix_best[0, 0]  # 最优风机转速
            optimal_backpressure = matrix_best[0, 1]  # 最优背压
            optimal_net_power = matrix_best[0, 2]  # 最优净功率
            
            logger.info("优化完成!")
            logger.info(f"最优风机转速: {optimal_fan_speed:.2f} rpm")
            logger.info(f"最优背压: {optimal_backpressure:.3f} kPa")
            logger.info(f"最优净功率: {optimal_net_power:.2f} MW")
            
            # 生成结果图表
            generate_optimization_plots(trace_list, logger)
            
            return {
                'optimal_fan_speed': optimal_fan_speed,
                'optimal_backpressure': optimal_backpressure,
                'optimal_net_power': optimal_net_power,
                'trace_list': trace_list,
                'optimization_time': need_time[0] if isinstance(need_time, list) else need_time
            }
        else:
            logger.error("优化算法返回空结果")
            return None
            
    except Exception as e:
        logger.error(f"优化过程中出现错误: {str(e)}")
        return None

def generate_optimization_plots(trace_list, logger):
    """生成优化结果图表"""
    try:
        if not trace_list or len(trace_list) == 0:
            logger.warning("没有优化轨迹数据，跳过图表生成")
            return
        
        # 创建结果目录
        os.makedirs("optimization_results", exist_ok=True)
        
        # 转换为numpy数组
        result_array = np.array(trace_list)
        
        # 创建图表
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
        
        # 1. 风机转速-净功率关系
        if result_array.shape[1] > 3:
            fan_speeds = result_array[:, 3]  # 风机转速
            net_powers = result_array[:, 1]  # 净功率
            ax1.plot(fan_speeds, net_powers, 'o-', linewidth=2, markersize=6)
            ax1.set_xlabel('风机转速 (rpm)')
            ax1.set_ylabel('净功率 (MW)')
            ax1.set_title('风机转速-净功率关系')
            ax1.grid(True, alpha=0.3)
        
        # 2. 背压-净功率关系
        if result_array.shape[1] > 5:
            backpressures = result_array[:, 5]  # 背压
            ax2.plot(backpressures, net_powers, 's-', linewidth=2, markersize=6, color='orange')
            ax2.set_xlabel('背压 (kPa)')
            ax2.set_ylabel('净功率 (MW)')
            ax2.set_title('背压-净功率关系')
            ax2.grid(True, alpha=0.3)
        
        # 3. 风机转速-背压关系
        if result_array.shape[1] > 5:
            ax3.plot(fan_speeds, backpressures, '^-', linewidth=2, markersize=6, color='green')
            ax3.set_xlabel('风机转速 (rpm)')
            ax3.set_ylabel('背压 (kPa)')
            ax3.set_title('风机转速-背压关系')
            ax3.grid(True, alpha=0.3)
        
        # 4. 负荷-净功率关系
        if result_array.shape[1] > 2:
            loads = result_array[:, 2]  # 负荷
            ax4.plot(loads, net_powers, 'd-', linewidth=2, markersize=6, color='red')
            ax4.set_xlabel('负荷 (MW)')
            ax4.set_ylabel('净功率 (MW)')
            ax4.set_title('负荷-净功率关系')
            ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig("optimization_results/opt_core3_optimization_results.png", dpi=300, bbox_inches='tight')
        plt.close()
        
        logger.info("优化结果图表已生成: optimization_results/opt_core3_optimization_results.png")
        
    except Exception as e:
        logger.error(f"生成图表时出现错误: {str(e)}")

def run_customer_style_simulation():
    """
    运行客户风格的仿真，完全按照main.py的格式
    """
    print("=" * 60)
    print("运行客户风格的仿真 (opt_core3)")
    print("=" * 60)
    
    # 设置日志
    logger = setup_logger()
    
    # 生成客户格式的数据
    need_time = [pd.Timestamp('2025-05-29 07:29:36')]
    
    # 生成基础数据
    base_data = generate_sample_data()
    
    # 生成最新一条数据
    latest_one_new = [base_data]
    
    # 生成历史数据（模拟客户的数据格式）
    latest_n_new = []
    for i in range(8):  # 生成8条历史数据
        varied_data = generate_varied_data(base_data, variation_factor=0.05)
        latest_n_new.append(varied_data)
    
    enough_his = True
    
    logger.info("数据准备完成:")
    logger.info(f"时间戳: {need_time}")
    logger.info(f"最新数据维度: {len(latest_one_new)}")
    logger.info(f"历史数据维度: {len(latest_n_new)}")
    logger.info(f"历史数据足够: {enough_his}")
    
    # 运行优化
    result = opt_core3(need_time, latest_one_new, latest_n_new, enough_his, logger)
    
    if result:
        print("\n优化结果:")
        print(f"最优风机转速: {result['optimal_fan_speed']:.2f} rpm")
        print(f"最优背压: {result['optimal_backpressure']:.3f} kPa")
        print(f"最优净功率: {result['optimal_net_power']:.2f} MW")
        print(f"优化时间: {result['optimization_time']}")
        print("\n结果已保存到 optimization_results/ 目录")
    else:
        print("优化失败")
    
    return result

def run_batch_simulation():
    """
    运行批量仿真，测试不同工况
    """
    print("=" * 60)
    print("运行批量仿真测试")
    print("=" * 60)
    
    logger = setup_logger()
    results = []
    
    # 测试不同负荷工况
    load_scenarios = [350, 400, 450, 500]  # MW
    
    for load in load_scenarios:
        logger.info(f"测试负荷: {load} MW")
        
        # 生成对应负荷的数据
        base_data = generate_sample_data()
        base_data[1] = load  # 设置负荷
        base_data[2] = load * 1.1  # 调整主汽流量
        
        need_time = [pd.Timestamp('2025-05-29 07:29:36')]
        latest_one_new = [base_data]
        
        # 生成历史数据
        latest_n_new = []
        for i in range(5):
            varied_data = generate_varied_data(base_data, variation_factor=0.03)
            latest_n_new.append(varied_data)
        
        # 运行优化
        result = opt_core3(need_time, latest_one_new, latest_n_new, True, logger)
        
        if result:
            result['load'] = load
            results.append(result)
        
        print(f"负荷 {load} MW 优化完成")
    
    # 生成批量结果报告
    if results:
        generate_batch_report(results, logger)
    
    return results

def generate_batch_report(results, logger):
    """生成批量仿真报告"""
    try:
        os.makedirs("optimization_results", exist_ok=True)
        
        # 创建对比图表
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
        
        loads = [r['load'] for r in results]
        fan_speeds = [r['optimal_fan_speed'] for r in results]
        backpressures = [r['optimal_backpressure'] for r in results]
        net_powers = [r['optimal_net_power'] for r in results]
        
        # 负荷-净功率关系
        ax1.plot(loads, net_powers, 'o-', linewidth=2, markersize=8)
        ax1.set_xlabel('机组负荷 (MW)')
        ax1.set_ylabel('最优净功率 (MW)')
        ax1.set_title('负荷-净功率关系')
        ax1.grid(True, alpha=0.3)
        
        # 负荷-背压关系
        ax2.plot(loads, backpressures, 's-', linewidth=2, markersize=8, color='orange')
        ax2.set_xlabel('机组负荷 (MW)')
        ax2.set_ylabel('最优背压 (kPa)')
        ax2.set_title('负荷-背压关系')
        ax2.grid(True, alpha=0.3)
        
        # 负荷-风机转速关系
        ax3.plot(loads, fan_speeds, '^-', linewidth=2, markersize=8, color='green')
        ax3.set_xlabel('机组负荷 (MW)')
        ax3.set_ylabel('最优风机转速 (rpm)')
        ax3.set_title('负荷-风机转速关系')
        ax3.grid(True, alpha=0.3)
        
        # 背压-净功率关系
        ax4.plot(backpressures, net_powers, 'd-', linewidth=2, markersize=8, color='red')
        ax4.set_xlabel('背压 (kPa)')
        ax4.set_ylabel('净功率 (MW)')
        ax4.set_title('背压-净功率关系')
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig("optimization_results/batch_optimization_results.png", dpi=300, bbox_inches='tight')
        plt.close()
        
        # 生成文本报告
        with open("optimization_results/batch_optimization_report.txt", "w", encoding="utf-8") as f:
            f.write("批量优化仿真结果报告\n")
            f.write("=" * 50 + "\n\n")
            
            for result in results:
                f.write(f"负荷: {result['load']} MW\n")
                f.write(f"最优风机转速: {result['optimal_fan_speed']:.2f} rpm\n")
                f.write(f"最优背压: {result['optimal_backpressure']:.3f} kPa\n")
                f.write(f"最优净功率: {result['optimal_net_power']:.2f} MW\n")
                f.write("-" * 30 + "\n")
        
        logger.info("批量仿真报告已生成")
        
    except Exception as e:
        logger.error(f"生成批量报告时出现错误: {str(e)}")

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1:
        if sys.argv[1] == "batch":
            # 批量仿真
            run_batch_simulation()
        else:
            # 客户风格仿真
            run_customer_style_simulation()
    else:
        # 默认运行客户风格仿真
        run_customer_style_simulation() 